using System;
using System.Diagnostics;
using System.Runtime.CompilerServices;
using System.Runtime.InteropServices;
using Unity.Collections;
using Unity.Collections.LowLevel.Unsafe;
using Unity.Mathematics;

namespace HTCollections.Unsafe
{
    /// <summary>
    /// 二维KD树
    /// </summary>
    [StructLayout(LayoutKind.Sequential)]
    [DebuggerDisplay("Length = {Length}, MaxLeafSize = {MaxLeafSize}, IsBuild = {IsBuild}")]
    [DebuggerTypeProxy(typeof(UnsafeKdTree2DDebugView<float3>))]
    public unsafe struct UnsafeKdTree2D : IDisposable
    {
        /// <summary>
        /// 最大叶子节点大小
        /// </summary>
        public int MaxLeafSize;

        [NativeDisableUnsafePtrRestriction] public UnsafeNode* NodesPtr;
        [NativeDisableUnsafePtrRestriction] public float3* PointsPtr; //float3(float2(pos), float(id))

        public int Length;
        public bool IsBuild;
        public Allocator Allocator;

        public UnsafeKdTree2D(int maxLeafSize = 5, Allocator allocator = Allocator.Persistent)
        {
            Allocator = allocator;
            MaxLeafSize = maxLeafSize;
            NodesPtr = null;
            PointsPtr = null;
            Length = 0;
            IsBuild = false;
        }

        public void Dispose()
        {
            if (!IsBuild) return;
            UnsafeUtil.Free(NodesPtr, Allocator);
            NodesPtr = null;
            PointsPtr = null;
            Length = 0;
            Allocator = Allocator.Invalid;
            IsBuild = false;
        }

        /// <summary>
        /// 构建KD树
        /// </summary>
        /// <param name="pointsPtr">坐标点集合指针</param>
        /// <param name="length">坐标点个数</param>
        public void Build(float3* pointsPtr, int length)
        {
            if (pointsPtr == null || length == 0) return;
            int pointLength = length;
            //PointsPtr = pointsPtr; //引用传递

            if (!IsBuild)
            {
                Length = pointLength * 2;
                NodesPtr = UnsafeUtil.Malloc<UnsafeNode>(Length, Allocator);
                PointsPtr = UnsafeUtil.Malloc<float3>(pointLength, Allocator); //内存拷贝
                UnsafeUtil.MemCpy(pointsPtr, PointsPtr, pointLength);
                IsBuild = true;
            }
            else
            {
                if (Length != pointLength * 2)
                {
                    UnsafeUtil.Free(NodesPtr, Allocator);
                    Length = pointLength * 2;
                    NodesPtr = UnsafeUtil.Malloc<UnsafeNode>(Length, Allocator);
                    UnsafeUtil.Free(PointsPtr, Allocator);
                    PointsPtr = UnsafeUtil.Malloc<float3>(pointLength, Allocator); //内存拷贝
                    UnsafeUtil.MemCpy(pointsPtr, PointsPtr, pointLength);
                }
                else
                {
                    UnsafeUtil.MemCpy(pointsPtr, PointsPtr, pointLength); //内存拷贝
                }
            }

            BuildRecursive(0, pointLength - 1, 0);
        }

        void BuildRecursive(int begin, int end, int nodeIndex)
        {
            UnsafeNode* nodePtr = &NodesPtr[nodeIndex];
            nodePtr->Index = nodeIndex;
            nodePtr->Begin = begin;
            nodePtr->End = end;
            nodePtr->MinX = float.MaxValue;
            nodePtr->MaxX = float.MinValue;
            nodePtr->MinY = float.MaxValue;
            nodePtr->MaxY = float.MinValue;
            nodePtr->LeftIndex = -1;
            nodePtr->RightIndex = -1;

            //计算当前节点的最小和最大坐标
            for (int i = begin; i <= end; i++)
            {
                NodesPtr[nodeIndex].MinX = Min(NodesPtr[nodeIndex].MinX, PointsPtr[i].x);
                NodesPtr[nodeIndex].MaxX = Max(NodesPtr[nodeIndex].MaxX, PointsPtr[i].x);
                NodesPtr[nodeIndex].MinY = Min(NodesPtr[nodeIndex].MinY, PointsPtr[i].y);
                NodesPtr[nodeIndex].MaxY = Max(NodesPtr[nodeIndex].MaxY, PointsPtr[i].y);
            }

            //如果当前节点的代理数量小于等于最大叶子节点大小，则不需要继续递归构建子树
            if (end - begin + 1 <= MaxLeafSize) return;

            int leftEnd = begin - 1;
            //计算当前节点的分割方向和分割值
            if (NodesPtr[nodeIndex].MaxX - NodesPtr[nodeIndex].MinX >
                NodesPtr[nodeIndex].MaxY - NodesPtr[nodeIndex].MinY) //沿竖直方向进行分割
            {
                float splitValue = (NodesPtr[nodeIndex].MaxX + NodesPtr[nodeIndex].MinX) / 2;
                //划分代理到左右子树
                for (int i = begin; i <= end; i++)
                {
                    if (PointsPtr[i].x <= splitValue) //左子树
                    {
                        leftEnd++;
                        Swap(i, leftEnd);
                    }
                }
            }
            else //沿水平方向进行分割
            {
                float splitValue = (NodesPtr[nodeIndex].MaxY + NodesPtr[nodeIndex].MinY) / 2;
                //划分代理到左右子树
                for (int i = begin; i <= end; i++)
                {
                    if (PointsPtr[i].y <= splitValue) //左子树
                    {
                        leftEnd++;
                        Swap(i, leftEnd);
                    }
                }
            }

            //如果左子树或右子树为空，则不需要继续递归构建子树
            if (leftEnd < begin || leftEnd >= end) return;

            //递归构建左子树和右子树
            NodesPtr[nodeIndex].LeftIndex = nodeIndex + 1;
            NodesPtr[nodeIndex].RightIndex = nodeIndex + 2 * (leftEnd - begin + 1);
            BuildRecursive(begin, leftEnd, NodesPtr[nodeIndex].LeftIndex);
            BuildRecursive(leftEnd + 1, end, NodesPtr[nodeIndex].RightIndex);
        }

        /// <summary>
        /// 搜索邻近点
        /// </summary>
        /// <param name="point">搜索中心位置坐标</param>
        /// <param name="range">搜索半径</param>
        /// <param name="results">搜索结果</param>
        public void QueryNeighborList(float3 point, float range, ref NativeList<float3> results)
        {
            if (Length == 0) return;
            QueryNeighborList(point, range * range, 0, ref results);
        }

        void QueryNeighborList(float3 point, float rangeSq, int nodeIndex, ref NativeList<float3> results)
        {
            //如果当前节点是叶子节点
            if (NodesPtr[nodeIndex].End - NodesPtr[nodeIndex].Begin + 1 <= MaxLeafSize)
            {
                for (int i = NodesPtr[nodeIndex].Begin; i <= NodesPtr[nodeIndex].End; i++)
                {
                    float disSq = DistanceSqrt(point.xy, PointsPtr[i].xy);
                    if (disSq < rangeSq)
                        results.Add(PointsPtr[i]);
                }
            }
            //如果当前节点不是叶子节点
            else
            {
                //计算指定代理到当前节点左子树和右子树的最短距离的平方
                //代理在左子树、右子树内部时最短距离视为0
                float disSq_left =
                    Sqrt(Max(0f, NodesPtr[NodesPtr[nodeIndex].LeftIndex].MinX - point.x)) +
                    Sqrt(Max(0f, point.x - NodesPtr[NodesPtr[nodeIndex].LeftIndex].MaxX)) +
                    Sqrt(Max(0f, NodesPtr[NodesPtr[nodeIndex].LeftIndex].MinY - point.y)) +
                    Sqrt(Max(0f, point.y - NodesPtr[NodesPtr[nodeIndex].LeftIndex].MaxY));
                float disSq_right =
                    Sqrt(Max(0f, NodesPtr[NodesPtr[nodeIndex].RightIndex].MinX - point.x)) +
                    Sqrt(Max(0f, point.x - NodesPtr[NodesPtr[nodeIndex].RightIndex].MaxX)) +
                    Sqrt(Max(0f, NodesPtr[NodesPtr[nodeIndex].RightIndex].MinY - point.y)) +
                    Sqrt(Max(0f, point.y - NodesPtr[NodesPtr[nodeIndex].RightIndex].MaxY));

                if (disSq_left < disSq_right) //如果左子树更近
                {
                    if (disSq_left < rangeSq) //如果左子树可能存在距离指定代理在一定范围内的邻居代理
                    {
                        QueryNeighborList(point, rangeSq, NodesPtr[nodeIndex].LeftIndex, ref results); //递归查询左子树
                        if (disSq_right < rangeSq) //如果右子树也可能存在距离指定代理在一定范围内的邻居代理
                            QueryNeighborList(point, rangeSq, NodesPtr[nodeIndex].RightIndex, ref results); //递归查询右子树
                    }
                }
                else //如果右子树更近
                {
                    if (disSq_right < rangeSq) //如果右子树可能存在距离指定代理在一定范围内的邻居代理
                    {
                        QueryNeighborList(point, rangeSq, NodesPtr[nodeIndex].RightIndex, ref results); //递归查询右子树
                        if (disSq_left < rangeSq) //如果左子树可能存在距离指定代理在一定范围内的邻居代理
                            QueryNeighborList(point, rangeSq, NodesPtr[nodeIndex].LeftIndex, ref results); //递归查询左子树
                    }
                }
            }
        }

        /// <summary>
        /// 搜索指定矩形范围内的点
        /// </summary>
        /// <param name="rect">指定矩形范围（minX, minY, maxX, maxY）</param>
        /// <param name="results">结果集合</param>
        public void QueryNeighborList(float4 rect, UnsafeList<float3>* results)
        {
            if (Length == 0) return;
            QueryNeighborList(rect, 0, results);
        }

        void QueryNeighborList(float4 rect, int nodeIndex, UnsafeList<float3>* results)
        {
            //如果当前节点是叶子节点
            var node = &NodesPtr[nodeIndex];
            if (node->End - node->Begin + 1 <= MaxLeafSize)
            {
                for (int i = node->Begin; i <= node->End; i++)
                {
                    if (Contain(PointsPtr[i], rect))
                        results->Add(PointsPtr[i]);
                }
            }
            //如果当前节点不是叶子节点
            else
            {
                //如果左子树可能存在矩形区域内的代理
                if (Contain(rect, &NodesPtr[node->LeftIndex]))
                {
                    QueryNeighborList(rect, node->LeftIndex, results); //递归查询左子树
                }

                //如果右子树可能存在矩形区域内的代理
                if (Contain(rect, &NodesPtr[node->RightIndex]))
                {
                    QueryNeighborList(rect, node->RightIndex, results); //递归查询右子树
                }
            }
        }

        /// <summary>
        /// 搜索指定范围内的最邻近点
        /// </summary>
        /// <param name="point">搜索中心位置坐标</param>
        /// <param name="range">搜索半径</param>
        /// <returns>坐标索引</returns>
        public int QueryNearest(float3 point, float range)
        {
            float rangeSq = float.MaxValue;
            int index = -1;
            QueryNearest(point, ref rangeSq, 0, ref index);
            if (rangeSq >= range * range) return -1;
            return index;
        }

        void QueryNearest(float3 point, ref float rangeSq, int nodeIndex, ref int index)
        {
            //如果当前节点是叶子节点
            if (NodesPtr[nodeIndex].End - NodesPtr[nodeIndex].Begin + 1 <= MaxLeafSize)
            {
                for (int i = NodesPtr[nodeIndex].Begin; i <= NodesPtr[nodeIndex].End; i++)
                {
                    float disSq = DistanceSqrt(point.xy, PointsPtr[i].xy);
                    if (disSq < rangeSq)
                    {
                        rangeSq = disSq;
                        index = i;
                    }
                }
            }
            //如果当前节点不是叶子节点
            else
            {
                //计算指定代理到当前节点左子树和右子树的最短距离的平方
                //代理在左子树、右子树内部时最短距离视为0
                float disSq_left =
                    Sqrt(Max(0f, NodesPtr[NodesPtr[nodeIndex].LeftIndex].MinX - point.x)) +
                    Sqrt(Max(0f, point.x - NodesPtr[NodesPtr[nodeIndex].LeftIndex].MaxX)) +
                    Sqrt(Max(0f, NodesPtr[NodesPtr[nodeIndex].LeftIndex].MinY - point.y)) +
                    Sqrt(Max(0f, point.y - NodesPtr[NodesPtr[nodeIndex].LeftIndex].MaxY));
                float disSq_right =
                    Sqrt(Max(0f, NodesPtr[NodesPtr[nodeIndex].RightIndex].MinX - point.x)) +
                    Sqrt(Max(0f, point.x - NodesPtr[NodesPtr[nodeIndex].RightIndex].MaxX)) +
                    Sqrt(Max(0f, NodesPtr[NodesPtr[nodeIndex].RightIndex].MinY - point.y)) +
                    Sqrt(Max(0f, point.y - NodesPtr[NodesPtr[nodeIndex].RightIndex].MaxY));

                if (disSq_left < disSq_right) //如果左子树更近
                {
                    if (disSq_left < rangeSq) //如果左子树可能存在距离指定代理在一定范围内的邻居代理
                    {
                        QueryNearest(point, ref rangeSq, NodesPtr[nodeIndex].LeftIndex, ref index); //递归查询左子树
                        if (disSq_right < rangeSq) //如果右子树也可能存在距离指定代理在一定范围内的邻居代理
                            QueryNearest(point, ref rangeSq, NodesPtr[nodeIndex].RightIndex, ref index); //递归查询右子树
                    }
                }
                else //如果右子树更近
                {
                    if (disSq_right < rangeSq) //如果右子树可能存在距离指定代理在一定范围内的邻居代理
                    {
                        QueryNearest(point, ref rangeSq, NodesPtr[nodeIndex].RightIndex, ref index); //递归查询右子树
                        if (disSq_left < rangeSq) //如果左子树可能存在距离指定代理在一定范围内的邻居代理
                            QueryNearest(point, ref rangeSq, NodesPtr[nodeIndex].LeftIndex, ref index); //递归查询左子树
                    }
                }
            }
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        void Swap(int a, int b)
        {
            var temp = PointsPtr[a];
            PointsPtr[a] = PointsPtr[b];
            PointsPtr[b] = temp;
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        float Sqrt(float a)
        {
            return a * a;
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        float DistanceSqrt(float2 a, float2 b)
        {
            return math.distancesq(a, b);
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        float Min(float a, float b)
        {
            if (a > b) return b;
            return a;
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        float Max(float a, float b)
        {
            if (a > b) return a;
            return b;
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        bool Contain(float3 point, float4 rect)
        {
            return point.x >= rect.x && point.x <= rect.z && point.y >= rect.y && point.y <= rect.w;
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        bool Contain(float4 rect, UnsafeNode* node)
        {
            if (node->MaxX < rect.x || node->MinX > rect.z) return false;
            if (node->MaxY < rect.y || node->MinY > rect.w) return false;
            return true;
        }

        public struct UnsafeNode : IEquatable<UnsafeNode>
        {
            public int Index; //节点Id
            public int Begin; //起始索引
            public int End; //终止索引
            public float MinX; //x轴上的最小值
            public float MaxX; //x轴上的最大值
            public float MinY; //y轴上的最小值
            public float MaxY; //y轴上的最大值
            public int LeftIndex; //左子树节点索引
            public int RightIndex; //右子树节点索引

            public UnsafeNode(int index, int begin, int end, float minX, float maxX, float minY, float maxY,
                int leftIndex,
                int rightIndex)
            {
                Index = index;
                Begin = begin;
                End = end;
                MinX = minX;
                MaxX = maxX;
                MinY = minY;
                MaxY = maxY;
                LeftIndex = leftIndex;
                RightIndex = rightIndex;
            }

            public bool Equals(UnsafeNode other)
            {
                return Index == other.Index && Begin == other.Begin && End == other.End && MinX.Equals(other.MinX) &&
                       MaxX.Equals(other.MaxX) && MinY.Equals(other.MinY) && MaxY.Equals(other.MaxY) &&
                       LeftIndex == other.LeftIndex && RightIndex == other.RightIndex;
            }

            public override bool Equals(object obj)
            {
                return obj is UnsafeNode other && Equals(other);
            }

            public override int GetHashCode()
            {
                return Index.GetHashCode();
            }
        }
    }

    sealed unsafe class UnsafeKdTree2DDebugView<T> where T : unmanaged
    {
        T* mPointsPtr;
        int mLength;

        public UnsafeKdTree2DDebugView(UnsafeKdTree2D array)
        {
            mPointsPtr = (T*)array.PointsPtr;
            mLength = array.Length / 2;
        }

        public T[] Items
        {
            get
            {
                T[] dst = new T[mLength];
                unsafe
                {
                    UnsafeUtil.CopyArray(mPointsPtr, 0, dst, 0, mLength);
                }

                return dst;
            }
        }
    }
}